package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/go-redis/redis/v8"
	"github.com/gorilla/websocket"
)

const idPathPrefix = "/intelligent-doc/id/"

var upgrader = websocket.Upgrader{
	// 允许所有CORS跨域请求。
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()

	logger := log.New(os.Stderr, "", log.LstdFlags)
	messageChan := make(chan string)
	redis_host := os.Getenv("REDIS_HOST")
	redis_port := os.Getenv("REDIS_PORT")
	redis_password := os.Getenv("REDIS_PASSWORD")
	redis_db := os.Getenv("REDIS_DB")

	if redis_host == "" {
		redis_host = "localhost"
	}

	if redis_port == "" {
		redis_port = "6379"
	}

	if redis_password == "" {
		redis_password = "difyai123456"
	}

	if redis_db == "" {
		redis_db = "0"
	}

	db_number, err := strconv.Atoi(redis_db)
	if err != nil {
		db_number = 0
	}

	// 创建一个新的 Redis 客户端
	rdb := redis.NewClient(&redis.Options{
		Addr:     redis_host + ":" + redis_port, // Redis 服务地址
		Password: redis_password,                // 密码
		DB:       db_number,                     // 使用默认数据库
	})

	go func() {
		for {
			_, message, err := conn.ReadMessage()
			if err != nil {
				logger.Println("read error:", err)
				close(messageChan)
				break
			}
			messageChan <- string(message) // 将消息发送到 channel 中
		}
	}()

	ctx := context.Background()
	id := strings.TrimPrefix(r.URL.Path, idPathPrefix)

	logger.Printf("id %v \n", id)
	count := 0
	timeout := 0
exit:
	for {
		select {
		case _, ok := <-messageChan:
			//用于检测客户端是否断开
			if !ok {
				break exit
			}
			//fmt.Println("received message:", msg)
		case <-time.After(time.Millisecond * 5):
			//等待是否会生成, 如果等待时间超过3分钟，则认为已经结束，可以等退出了
			//因为生成时间可能很长
			//暂时由客户端来断开
			/*
				if timeout > 3000 {
					time.Sleep(100 * time.Millisecond)
				}
			*/

			// 正常处理消息...
			key := id + fmt.Sprintf("-%05d", count)
			val, err := rdb.Get(ctx, key).Result()
			if err == redis.Nil {
				if timeout == 0 {
					logger.Printf("next key %v does not generate, wait ...", key)
				}
				time.Sleep(300 * time.Millisecond)
				timeout++
			} else if err != nil {
				logger.Printf("key %v \n", key)
				logger.Printf("redis read err %v", err)
				break exit
			} else {
				logger.Printf("key %v \n", key)
				logger.Println("value:", val)
				count++
				timeout = 0
				if err := conn.WriteMessage(websocket.TextMessage, []byte(val)); err != nil {
					if err == websocket.ErrCloseSent {
						break exit
					}
					logger.Println("websocket write err: ", err)
				}
			}
		}
	}

	logger.Println("client ", id, " disconnected normally")
}

func main() {
	logger := log.New(os.Stderr, "", log.LstdFlags)
	logger.Println("websocket server start ...")

	http.HandleFunc(idPathPrefix, wsHandler)
	log.Fatal(http.ListenAndServe(":9527", nil))
}
